import torch
import torch.nn as nn
from torchvision import transforms, datasets
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt
import time

from Lesson.CNN.lession12_CNN import CNN

x=torch.randn(8, 3, 80, 80)
conv=nn.Conv2d(3, 6, kernel_size=(5,5))
print(conv(x).shape)

# pool = nn.AvgPool2d(2, 3)
# print(pool(x).shape)

pool = nn.AdaptiveAvgPool2d(13)
# 归一化
bn1=nn.BatchNorm1d(3)
a=torch.randn(2, 3)

# pipline=transforms.Compose([transforms.ToTensor(), transforms.Normalize([0.5], [0.5])])
# train_dataset = datasets.CIFAR10('CIFAR10_data', train=True, transform=pipline,download=True)

p = transforms.Compose([transforms.ToTensor(), transforms.Normalize([0.5], [0.5])])
train = datasets.MNIST('../data', transform=p, train=True, download=True)
train_loader = DataLoader(train, batch_size=16, shuffle=True)

net = CNN()
loss_fun = nn.CrossEntropyLoss()
optim = torch.optim.SGD(net.parameters(), lr=0.01)

losses = []
acces = []
start_time = time.time()
num_epoch = 1

# 取出每一行中最大值下标
a = torch.randn(3, 4)
val, index = a.max(dim=1)
idx = a.argmax(1)

start = time.time()
for epoch in range(num_epoch):
    train_loss = 0.0
    train_acc = 0.0
    for img, label in train_loader:
        img = img.view(img.shape[0], -1)
        out = net(img)
        optim.zero_grad()
        loss = loss_fun(out, label)
        loss.backward()
        optim.step()
        pred = out.argmax(dim=1)
        num_correct = (pred==label).sum().item()
        acc = num_correct/img.shape[0]
        train_acc += acc
        train_loss += loss.item()
    losses.append(train_loss/len(train_loader))
    acces.append(train_acc/len(train_loader))
    print(f'Epoch:{epoch + 1}, train_loss:{train_loss / len(train_loader):.4f}, train_acc:{train_acc / len(train_loader):.4f}')

end = time.time()

print(f'训练时长:{end-start}')

plt.figure(figsize=(10, 4))
plt.subplot(121)
plt.ylabel('train_loss')
plt.xlabel('Epoch')
plt.plot(range(num_epoch), losses)
plt.subplot(122)
plt.ylabel('train_acces')
plt.xlabel('Epoch')
plt.plot(range(num_epoch), acces, r='r')
plt.show()

# GPU
# device=torch.device('cuda' if )

#读取一个图片，自己的数据集
dog_data=datasets.ImageFolder(r'')
